//
// Created by hamlet on 25-2-7.
//

#ifndef FUNCTION_H
#define FUNCTION_H

// ----------------------------------------------------------------------------
// all the individual layers' forward passes
// B = batch_size, T = sequence_length, C = channels, V = vocab_size

/**
 * @brief 编码器前向传播
 * @param out 输出张量
 * @param inp 输入token序列
 * @param wte token嵌入权重
 * @param wpe 位置嵌入权重
 * @param B 批次大小
 * @param T 序列长度
 * @param C 嵌入维度
 */
void encoder_forward(float* out,
                     const int* inp, const float* wte, const float* wpe,
                     int B, int T, int C);

/**
 * @brief 层归一化前向传播
 * @param out 归一化后的输出
 * @param mean 均值
 * @param rstd 标准差的倒数
 * @param inp 输入张量
 * @param weight 缩放参数
 * @param bias 偏置参数
 * @param B 批次大小
 * @param T 序列长度
 * @param C 特征维度
 */
void layernorm_forward(float* out, float* mean, float* rstd,
                       const float* inp, const float* weight, const float* bias,
                       int B, int T, int C);

/**
 * @brief 矩阵乘法前向传播
 * @param out 输出张量
 * @param inp 输入张量
 * @param weight 权重矩阵
 * @param bias 偏置向量
 * @param B 批次大小
 * @param T 序列长度
 * @param C 输入特征维度
 * @param OC 输出特征维度
 */
void matmul_forward(float* out,
                    const float* inp, const float* weight, const float* bias,
                    int B, int T, int C, int OC);

/**
 * @brief 注意力机制前向传播
 * @param out 注意力输出
 * @param preatt 注意力计算的中间结果
 * @param att 注意力权重
 * @param inp 输入张量
 * @param B 批次大小
 * @param T 序列长度
 * @param C 特征维度
 * @param NH 注意力头数
 */
void attention_forward(float* out, float* preatt, float* att,
                       const float* inp,
                       int B, int T, int C, int NH);

/**
 * @brief GELU激活函数前向传播
 * @param out 输出张量
 * @param inp 输入张量
 * @param N 张量元素总数
 */
void gelu_forward(float* out, const float* inp, int N);

/**
 * @brief 残差连接前向传播
 * @param out 输出张量
 * @param inp1 第一个输入张量
 * @param inp2 第二个输入张量
 * @param N 张量元素总数
 */
void residual_forward(float* out, const float* inp1, const float* inp2, int N);

/**
 * @brief Softmax前向传播
 * @param probs 输出概率分布
 * @param logits 输入logits
 * @param B 批次大小
 * @param T 序列长度
 * @param V 词表大小
 * @param Vp 填充后的词表大小(用于内存对齐)
 */
void softmax_forward(float* probs, const float* logits, int B, int T, int V, int Vp);

#endif //FUNCTION_H
